/*****************************************************************
 *   Licensed to the Apache Software Foundation (ASF) under one
 *  or more contributor license agreements.  See the NOTICE file
 *  distributed with this work for additional information
 *  regarding copyright ownership.  The ASF licenses this file
 *  to you under the Apache License, Version 2.0 (the
 *  "License"); you may not use this file except in compliance
 *  with the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an
 *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 *  KIND, either express or implied.  See the License for the
 *  specific language governing permissions and limitations
 *  under the License.
 ****************************************************************/
package org.epiphron.wam;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;

import org.apache.bcel.Constants;
import org.apache.bcel.Repository;
import org.apache.bcel.classfile.Code;
import org.apache.bcel.classfile.ConstantCP;
import org.apache.bcel.classfile.ConstantNameAndType;
import org.apache.bcel.classfile.ConstantPool;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.LineNumberTable;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.CPInstruction;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.Instruction;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.LDC;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Call graph of functions in the Web Application
 */
public class CallGraph<T1> {
	static final Logger logger = LoggerFactory.getLogger(CallGraph.class);
	
	/**
	 * Descriptor of a node, which is essentially a method
	 */
	public static class MethodNode<T> {
		/**
		 * Method's id
		 */
		String id;
		
		/**
		 * Linked nodes
		 */
		List<MethodNode<T>> linked;
		
		/**
		 * Invoked operations
		 */
		Collection<T> operations;
		
		MethodNode(String id) {
			linked = new ArrayList<MethodNode<T>>();
			operations = new ArrayList<T>();
			
			this.id = id;
		}
		
		private void addNode(MethodNode<T> node) {
			linked.add(node);
		}
		
		private List<MethodNode<T>> getLinked() {
			return linked;
		}
		
		private void addOperation(T op) {
			operations.add(op);
		}
		
		public Collection<T> getOperations() {
			return operations;
		}
		
		@Override
		public int hashCode() {
			return id.hashCode();
		}
		
		@Override
		public boolean equals(Object obj) {
			return obj instanceof MethodNode && ((MethodNode<?>) obj).id.equals(id);
		}
	}
	
	Map<String, MethodNode<T1>> methodCache;
	
	public CallGraph() {
		methodCache = new HashMap<String, MethodNode<T1>>();
	}
	
	/**
	 * Builds, if necessary, and returns list of called methods,
	 * starting with a given method
	 * @throws ClassNotFoundException 
	 */
	public Collection<MethodNode<T1>> buildCalledMethods(java.lang.reflect.Method method,
			IOperationFactory<T1> opFactory) throws ClassNotFoundException {
		Collection<MethodNode<T1>> nodes = new LinkedHashSet<MethodNode<T1>>();
		
		JavaClass clazz = Repository.getRepository().loadClass(method.getDeclaringClass());
		Method methodDesc = clazz.getMethod(method);
		
		MethodNode<T1> rootNode = crawl(method.getDeclaringClass().getName(), methodDesc, opFactory);
		collect(nodes, rootNode);
		return nodes;
	}
	
	private String buildId(String className, Method methodDesc) {
		return className + "." + methodDesc.getName() + " " + methodDesc.getSignature();
	}

	private MethodNode<T1> crawl(String className, Method methodDesc, IOperationFactory<T1> opFactory) {
		String id = buildId(className, methodDesc);
		
		MethodNode<T1> node;
		if (methodCache.containsKey(id)) {
			node = (MethodNode<T1>) methodCache.get(id);
		}
		else {
			node = new MethodNode<T1>(id);
			
			// adding to cache right away
			methodCache.put(id, node);
			
			// fill in
			parse(className, methodDesc, node, opFactory);
		}
		
		return node;
	}

	private void parse(String className, Method methodDesc, MethodNode<T1> node, IOperationFactory<T1> opFactory) {
		ConstantPool pool = methodDesc.getConstantPool();
		LineNumberTable lineNumbers = methodDesc.getLineNumberTable();
		
		Code code = methodDesc.getCode();
		if (code != null) {
			// reading all invoke instructions
			InstructionList instList = new InstructionList(code.getCode());
			
			Instruction[] instructions = instList.getInstructions();
			for (int i = 0; i < instructions.length; i++) {
				Instruction instruction = instructions[i];
				
				byte type = Constants.CONSTANT_Methodref;
				
				switch (instruction.getOpcode()) {
				case Constants.INVOKEINTERFACE:
					type = Constants.CONSTANT_InterfaceMethodref;
					// fall through
				case Constants.INVOKESTATIC:
				case Constants.INVOKEVIRTUAL:
				case Constants.INVOKENONVIRTUAL:
					// getting real value from constant pool
					CPInstruction cpi = (CPInstruction) instruction;
					int index = cpi.getIndex();
					
					ConstantCP constant = (ConstantCP) pool.getConstant(index, type);
					
					String instrClass = constant.getClass(pool);
					
					ConstantNameAndType nameAndType = (ConstantNameAndType)
							pool.getConstant(constant.getNameAndTypeIndex());
					String methodName = pool.constantToString(nameAndType.getNameIndex(), Constants.CONSTANT_Utf8);
					
					IOperationDescription opDesc = Operations.find(instrClass, methodName);
					if (opDesc != null) {
						Object previousConstant = null;
						if (opDesc.acceptsParameter() && i > 0) {
							Instruction prev = instructions[i - 1];
							if (prev.getOpcode() == Constants.LDC) {
								previousConstant = ((LDC) prev).getValue(new ConstantPoolGen(pool));
							}
						}
						
						logger.info("Found operation " + opDesc.getName() + "(" + (previousConstant != null ? previousConstant : "")
								+ ") in method " + className + "." + methodDesc.getName());
						
						int line = -1;
						if (lineNumbers != null) {
							line = lineNumbers.getSourceLine(instList.getInstructionPositions()[i]);
						}
						
						T1 op = opFactory.buildOperation(opDesc,
							className, methodDesc.getName(), methodDesc.getSignature(),
							line, previousConstant);
						node.addOperation(op);
					}
					else if (!instrClass.startsWith("java.") && !instrClass.startsWith("javax.")) { // optimizing
						String fullSignature = pool.constantToString(constant);
						String signature = pool.constantToString(nameAndType.getSignatureIndex(), Constants.CONSTANT_Utf8);
						
						// check if we don't need to load the class but can instead get straight from cache
						if (methodCache.containsKey(fullSignature)) {
							node.addNode(methodCache.get(fullSignature));
						}
						else {
							JavaClass clazz = Repository.lookupClass(instrClass);
							if (clazz != null) {
								Method m = null;
								for (Method meth : clazz.getMethods()) {
									if (meth.getName().equals(methodName) && meth.getSignature().equals(signature)) {
										m = meth;
										break;
									}
								}
								
								if (m != null) {
									MethodNode<T1> referencedNode = crawl(instrClass, m, opFactory);
									node.addNode(referencedNode);
								}
							}
							else {
								logger.warn("Referenced class not found: " + instrClass);
							}
						}
					}
				}
			}
		}
	}
	
	private void collect(Collection<MethodNode<T1>> nodes, MethodNode<T1> rootNode) {
		if (!nodes.contains(rootNode)) {
			nodes.add(rootNode);
			
			for (MethodNode<T1> linked : rootNode.getLinked()) {
				collect(nodes, linked);
			}
		}
	}
}
